package com.migration.model.migration.scripts;

import com.migration.model.migration.AbstractMultiThreadMigrationScript;
import com.migration.model.migration.CustomExecutableMigrationScript;
import com.migration.model.newdb.RiskChainEffect;
import com.migration.model.olddb.*;
import com.migration.service.RiskChainDetailService;
import com.migration.service.RiskChainEffectService;
import jakarta.persistence.EntityManager;
import jakarta.persistence.EntityManagerFactory;
import jakarta.persistence.criteria.CriteriaBuilder;
import jakarta.persistence.criteria.CriteriaQuery;
import jakarta.persistence.criteria.Root;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.stereotype.Component;

import javax.sql.DataSource;
import java.util.*;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 事故链-影响表(risk_chain_effect)及其关联表迁移脚本
 */
@Slf4j
@Component
public class A025_RiskChainEffectMigrationScript extends AbstractMultiThreadMigrationScript implements CustomExecutableMigrationScript {

    @Autowired
    @Qualifier("oldEntityManagerFactory")
    private EntityManagerFactory oldEntityManagerFactory;

    @Autowired
    @Qualifier("newEntityManagerFactory")
    private EntityManagerFactory newEntityManagerFactory;

    @Autowired
    private RiskChainEffectService riskChainEffectService;

    @Autowired
    private RiskChainDetailService riskChainDetailService;

    /**
     * 构造函数，初始化迁移配置
     */
    public A025_RiskChainEffectMigrationScript() {
        super(
                "risk_chain_effect",
                "事故链-影响表及关联表数据迁移",
                "ra_influence",
                "risk_chain_effect",
                25
        );
    }

    @Override
    protected String generateMigrationSql() {
        return "";
    }

    @Override
    public int executeCustomMigration(DataSource dataSource) {
        // 清空目标表
        JdbcTemplate jdbcTemplate = new JdbcTemplate(dataSource);
        jdbcTemplate.execute(getTruncateSql()); // 清空影响表
        jdbcTemplate.execute("TRUNCATE TABLE risk_chain_effect_chemical;");
        jdbcTemplate.execute("TRUNCATE TABLE risk_chain_effect_deviation;");
        jdbcTemplate.execute("TRUNCATE TABLE risk_chain_effect_next_device;");
        jdbcTemplate.execute("TRUNCATE TABLE risk_chain_effect_protection;");
        jdbcTemplate.execute("TRUNCATE TABLE risk_chain_effect_suggestion;");

        // 执行自定义迁移
        return executeWithEntityProcessingMultiThread(oldEntityManagerFactory, newEntityManagerFactory);
    }

    /**
     * 构建reasonId到RaDeviateReason的映射
     *
     * @param oldEntityManager 旧数据库实体管理器
     * @param reasonIds        原因ID集合
     * @return 原因ID到原因实体的映射
     */
    private Map<Long, RaDeviateReason> buildReasonMap(EntityManager oldEntityManager, Set<Long> reasonIds) {
        Map<Long, RaDeviateReason> reasonMap = new HashMap<>();
        if (!reasonIds.isEmpty()) {
            CriteriaBuilder cb = oldEntityManager.getCriteriaBuilder();
            CriteriaQuery<RaDeviateReason> cq = cb.createQuery(RaDeviateReason.class);
            Root<RaDeviateReason> root = cq.from(RaDeviateReason.class);
            cq.where(root.get("id").in(reasonIds));
            List<RaDeviateReason> reasonList = oldEntityManager.createQuery(cq).getResultList();

            for (RaDeviateReason reason : reasonList) {
                reasonMap.put(reason.getId(), reason);
            }
        }
        return reasonMap;
    }

    protected int executeWithEntityProcessingMultiThread(EntityManagerFactory oldEntityManagerFactory, EntityManagerFactory newEntityManagerFactory) {
        EntityManager oldEntityManager = oldEntityManagerFactory.createEntityManager();
        AtomicInteger rowsProcessed = createAtomicCounter();

        // 创建统计器
        AtomicInteger effectRows = new AtomicInteger(0);
        AtomicInteger effectRelationRows = new AtomicInteger(0);

        try {
            // 1. 获取所有事故链详情数据
            List<RaFaultChainsDetail> detailList = riskChainDetailService.getAllFaultChainsDetails(oldEntityManager);

            // 2. 获取所有事故链措施关系
            List<RaFaultChainsMeasure> measureList = riskChainDetailService.getAllFaultChainsMeasures(oldEntityManager);

            // 3. 从详情表中获取所有原因ID，用于构建reasonMap
            Set<Long> detailReasonIds = new HashSet<>();
            for (RaFaultChainsDetail detail : detailList) {
                if (detail.getReasonId() != null) {
                    detailReasonIds.add(detail.getReasonId());
                }
            }

            // 4. 构建reasonId到RaDeviateReason的映射
            Map<Long, RaDeviateReason> reasonMap = buildReasonMap(oldEntityManager, detailReasonIds);

            // 5. 获取所有影响数据
            List<RaInfluence> influenceList = riskChainEffectService.getAllInfluences(oldEntityManager);

            // 6. 构建影响节点ID映射
            Map<Long, RaInfluence> influenceMap = riskChainEffectService.buildInfluenceMap(influenceList);

            // 7. 构建源ID与影响节点ID的映射
            Map<Long, Set<Long>> sourceIdToInfluenceIds = riskChainEffectService.buildSourceIdToInfluenceIdsMap(
                    measureList,
                    detailList
            );

            // 8. 获取所有影响关联表数据
            List<RaInfluenceLinked> influenceLinkedList = riskChainEffectService.getAllInfluenceLinked(oldEntityManager);

            /*********************************************** 多线程迁移数据 **************************************************/
            // 创建线程池和计数器
            ExecutorService executorService = createExecutorService();
            CountDownLatch latch = createCountDownLatch(1);

            // 迁移影响节点数据及关联数据
            executorService.submit(() -> {
                try {
                    EntityManager newEntityManager = newEntityManagerFactory.createEntityManager();
                    try {
                        newEntityManager.getTransaction().begin();

                        // 批次迁移影响节点数据
                        List<RiskChainEffect> effectList = riskChainEffectService.migrateEffectData(
                            oldEntityManager,
                            newEntityManager,
                            influenceMap,
                            sourceIdToInfluenceIds,
                            detailList,
                            reasonMap
                        );

                        int size = effectList.size();
                        effectRows.addAndGet(size);
                        rowsProcessed.addAndGet(size);

                        log.info("事故链-影响表(risk_chain_effect)迁移完成，共迁移 {} 行数据", effectRows.get());

                        // 迁移关联数据
                        int relationRows = riskChainEffectService.migrateEffectRelationData(
                                oldEntityManager,
                                newEntityManager,
                                effectList,
                                influenceLinkedList
                        );

                        effectRelationRows.addAndGet(relationRows);
                        rowsProcessed.addAndGet(relationRows);

                        log.info("事故链-影响关联表(risk_chain_effect_*)迁移完成，共迁移 {} 行数据", effectRelationRows.get());

                        newEntityManager.getTransaction().commit();
                    } catch (Exception e) {
                        if (newEntityManager.getTransaction().isActive()) {
                            newEntityManager.getTransaction().rollback();
                        }
                        throw e;
                    } finally {
                        newEntityManager.close();
                        latch.countDown();
                    }
                } catch (Exception e) {
                    log.error("事故链-影响表及关联表迁移异常", e);
                    throw new RuntimeException(e);
                }
            });

            // 等待所有任务完成
            waitForCompletion(latch, executorService);

            // 打印总结日志
            log.info("事故链-影响表迁移总结 ========================");
            log.info("事故链-影响表(risk_chain_effect)：{} 行", effectRows.get());
            log.info("事故链-影响关联表(risk_chain_effect_*)：{} 行", effectRelationRows.get());
            log.info("总计迁移数据：{} 行", rowsProcessed.get());

        } catch (Exception e) {
            log.error("事故链-影响表及关联表迁移异常", e);
            throw e;
        } finally {
            oldEntityManager.close();
        }

        return rowsProcessed.get();
    }
}
